"""
实现分词
"""
import os
from torch.utils.data import DataLoader, Dataset
import re
from lib import ws
import torch
import lib


base_data_path = r"F:\virtual_environment\data\aclImdb_v1\aclImdb"


# 将数据转化为token
def tokens(text):
    # 准备要分割的字符
    filters = ['\.', '\t', '\n', '\x97', '\x96', '#', '$', '%', '&']

    # 将文本中的HTML标签删除
    text = re.sub("<.*?>", " ", text)

    # 将要分割的字符删除
    text = re.sub("|".join(filters), " ", text)

    # 返回文本tokens
    return [i.strip().lower() for i in text.split()]


# 获取数据
class ImDb(Dataset):
    def __init__(self, mode):
        # super().__init__()  # 方法一
        super(ImDb, self).__init__()
        if mode == True:
            # os.path.join的第一个参数不支持'+ 路径'
            data_path = [os.path.join(base_data_path, i) for i in ["train/pos", "train/neg"]]
        else:
            data_path = [os.path.join(base_data_path, i) for i in ["test/pos", "test/neg"]]

        self.text_file_list = [os.path.join(path, i) for path in data_path for i in os.listdir(path) if i.endswith(".txt")]
        # self.text_file_list = []
        # for i in data_path:
        #     self.text_file_list.extend([os.path.join(i, j) for j in os.listdir(i)])

    def __getitem__(self, index):
        file_path = self.text_file_list[index]
        # print(file_path)
        # 读取数据
        data = open(file_path, mode='r', encoding="UTF-8").read()
        # 获取标签
        # os.path.basename(now_path)获取当前路径的最后一个文件，比如F:/A/b.txt， 获取的是b.txt
        score = int(os.path.basename(file_path).split("_")[-1].split(".")[0])
        label = 0 if score < 5 else 1
        return tokens(data), label

    def __len__(self):
        return len(self.text_file_list)


def collate_fn(batch):
    content, label = zip(*batch)
    # max_len每个句子保留多少个词
    content = [ws.word_to_num_transform(words, max_len=lib.seq_len) for words in content]
    content = torch.LongTensor(content)  # 进行embedding操作的数据必须是LongTensor类型的
    label = torch.LongTensor(label)
    return content, label


def get_dataloader(train=True, batch_size=lib.train_batch_size):
    dataset = ImDb(train)
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True, collate_fn=collate_fn)
    return data_loader


if __name__ == '__main__':
    for idx, (target, label) in enumerate(get_dataloader()):
        print("idx=", idx)
        print("label=", label)
        print("target=", target)
        print(label.size())
        print(target.size())
        break
    # ImDb()[0]
    # print(ImDb()[0])
